1899G - Unusual Entertainment - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define int ll
typedef pair<int,int> pii;
#define X first
#define Y second
#define pb push_back
#define All(a) a.begin(), a.end()
#define SZ(a) ((int)a.size())
#define endl '\n'
#define yn(ok) cout << (ok ? "Yes" : "No") << endl;
void db() {cout << '\n';}
template <typename T, typename ...U> void db(T a, U ...b) {
cout << a << ' ', db(b...);}

const ll mod = 1e9+7;
const ll inf = 1LL << 60;
const int NN = 2e5+5;

int n, m, k;
vector<int> g[NN];
set<int> ch[NN];
vector<int> id, chnum;
vector<tuple<int, int, int>> q[NN];
vector<int> chnow, ans;

void dfs(int k, int p){
    chnum[k] = 1;
    chnow[k] = k;
    int mx = 0, mxch = k;
    for(int nt : g[k]){
        if(nt != p){
            dfs(nt, k);
//            ch[k].insert(id[nt]);
            if(chnum[nt] > mx){
                mx = chnum[nt], mxch = nt;
            } 
            chnum[k] += chnum[nt];
        } 
    } 
    chnow[k] = chnow[mxch];
    ch[chnow[k]].insert(id[k]);
    for(int nt : g[k]){
        if(nt != p && nt != mxch){
            for(int i : ch[chnow[nt]])
                ch[chnow[k]].insert(i);
        } 
    } 
//    db(k, ":", chnow[k]);
//    for(int i : ch[chnow[k]]) cout << i << ' '; cout << endl;
    for(auto[l, r, ansid] : q[k]){
        auto it = ch[chnow[k]].lower_bound(l);
        if(it != ch[chnow[k]].end() && *it <= r) ans[ansid] = 1;
    } 
} 

void solve(){
    cin >> n >> m;
    id.resize(n + 1);
    chnum.resize(n + 1);
    chnow.resize(n + 1);
    ans.assign(m, 0);
    for(int i = 0; i < n - 1; i++){
        int x, y; cin >> x >> y;
        g[x].pb(y);
        g[y].pb(x);
    } 
    vector<int> p(n + 1);
    for(int i = 1; i <= n; i++) cin >> p[i], id[p[i]] = i;
    for(int i = 0; i < m; i++){
        int x, y, z;
        cin >> x >> y >> z;
        q[z].pb({x, y, i});
    } 
    dfs(1, 0);
    for(int i : ans) yn(i);
    for(int i = 0; i <= n; i++) g[i].clear(), ch[i].clear();
    for(int i = 0; i <= m; i++) q[i].clear();
}

signed main(){
    ios_base::sync_with_stdio(false); cin.tie(0);
    int t = 1;
    cin >> t;
    while(t--)
        solve();
    return 0;
}